package edu.buaa.energy.client;

import com.google.common.util.concurrent.ListenableFuture;
import edu.buaa.common.client.PostgreSqlExecutorClient;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.energy.EnergyDataTxGenerator;
import edu.buaa.energy.transaction.*;
import edu.buaa.utils.Pair;
import edu.buaa.utils.Triple;

import java.sql.*;
import java.util.*;
import java.util.concurrent.*;

public class PostgreSQLClient extends PostgreSqlExecutorClient {

    public PostgreSQLClient(String serverHost, int parallelCnt) throws Exception {
        super(serverHost, "energy", parallelCnt);
    }

    @Override
    protected List<String> createTables() {
        return Arrays.asList(
                "CREATE table IF NOT EXISTS energy_node(id BIGSERIAL PRIMARY KEY, name varchar(255), country varchar(255), volt int, lat float, lon float)",
                "CREATE table IF NOT EXISTS electric_line(r_start int, r_end int, x float, y float, r_limit int, r_length float)",
                "CREATE table IF NOT EXISTS temporal_status(node_id int, t int, load_signal FLOAT, " +
                    "solar_cosmo_p FLOAT, solar_cosmo_u FLOAT, solar_ecmwf_p FLOAT, solar_ecmwf_u FLOAT, " +
                    "wind_cosmo_p FLOAT, wind_cosmo_u FLOAT, wind_ecmwf_p FLOAT, wind_ecmwf_u FLOAT, PRIMARY KEY (node_id,t))",
                "create index IF NOT EXISTS ind_id_t on temporal_status(node_id, t)"
        );
    }

    @Override
    public ListenableFuture<ServerResponse> execute(AbstractTransaction tx) throws InterruptedException {
        switch (tx.getTxType()) {
            case tx_import_static_data:
                return this.submit(execute((ImportStaticDataTx) tx));
            case tx_import_temporal_data:
                return this.submit(execute((ImportTemporalDataTx) tx), tx.getSection());
            case tx_query_snapshot:
                return this.submit(execute((SnapshotQueryTx) tx));
            case tx_update_temporal_data:
                return this.submit(execute((UpdateTemporalDataTx) tx));
            case tx_query_snapshot_aggr_max:
                return this.submit(execute((SnapshotAggrMaxTx) tx));
            case tx_query_snapshot_aggr_duration:
                return this.submit(execute((SnapshotAggrDurationTx) tx));
            case tx_query_entity_history:
                return this.submit(execute((EntityHistoryTx) tx));
            case tx_query_road_by_temporal_condition:
                return this.submit(execute((EntityTemporalConditionTx) tx));
            default:
                throw new UnsupportedOperationException();
        }
    }

    private Req execute(ImportStaticDataTx tx) {
        return conn -> {
            try (PreparedStatement stat1 = conn.prepareStatement("INSERT INTO energy_node VALUES (?,?,?,?,?,?)");
                 PreparedStatement stat2 = conn.prepareStatement("INSERT INTO electric_line VALUES (?, ?, ?, ?, ?, ?)")) {
                conn.setAutoCommit(false);
                for (ImportStaticDataTx.EnergyNode p : tx.getNodes()) {
                    stat1.setInt(1, p.getRawId());
                    stat1.setString(2, p.getName());
                    stat1.setString(3, p.getCountry());
                    stat1.setInt(4, p.getVoltage());
                    stat1.setFloat(5, p.getLatitude());
                    stat1.setFloat(6, p.getLongitude());
                    stat1.addBatch();
                }
                stat1.executeBatch();
                for (ImportStaticDataTx.EnergyRel r : tx.getRels()) {
                    stat2.setInt(1, r.getFromRawId());
                    stat2.setInt(2, r.getToRawId());
                    stat2.setFloat(3, r.getX());
                    stat2.setFloat(4, r.getY());
                    stat2.setInt(5, r.getLimit());
                    stat2.setFloat(6, r.getLength());
                    stat2.addBatch();
                }
                stat2.executeBatch();
                conn.commit();
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(ImportTemporalDataTx tx) {
        return conn -> {
            try (PreparedStatement stmt = conn.prepareStatement("insert into temporal_status values(?,?,?,?,?,?,?,?,?,?,?)")) {
                conn.setAutoCommit(false);
                for (ImportTemporalDataTx.EnergyStatus s : tx.data) {
                    stmt.setLong(1, s.getNodeId());
                    stmt.setInt(2, s.getTime());
                    List<Float> ss = s.getStatus();
                    for(int i=3; i<3+ss.size();i++){
                        stmt.setFloat(i, ss.get(i-3));
                    }
                    stmt.addBatch();
                }
                stmt.executeBatch();
                conn.commit();
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(UpdateTemporalDataTx tx) {
        return conn -> {
            conn.setAutoCommit(false);
            int nodeId = tx.getEntityId();
            String[] props = EnergyDataTxGenerator.properties;
            String propStr = String.join(",", props);
            try (PreparedStatement stmt1 = conn.prepareStatement("select node_id, max(t) as max_t from temporal_status where node_id=? and t<=? group by node_id");
                 PreparedStatement stmt2 = conn.prepareStatement("select "+propStr+" from temporal_status where node_id=? and t=?");
                 PreparedStatement stmt3 = conn.prepareStatement("insert into temporal_status(t, node_id, "+propStr+") values(?,?,?,?,?,?,?,?,?,?,?) " +
                         "on CONFLICT (node_id, t) DO NOTHING");
                 PreparedStatement stmt4 = conn.prepareStatement("delete from temporal_status where node_id=? and t>=? and t<=?")) {
                stmt1.setInt(1, nodeId);
                stmt1.setInt(2, tx.getEndTime());
                ResultSet rs = stmt1.executeQuery();

                // insert into table from t1 + 1 with the old value.
                if (rs.next()) {
                    stmt2.setInt(1, rs.getInt("node_id"));
                    stmt2.setInt(2, rs.getInt("max_t"));
                    ResultSet r = stmt2.executeQuery();
                    if (r.next()) {
                        stmt3.setInt(1, tx.getEndTime() + 1);
                        stmt3.setInt(2, nodeId);
                        int j = 3;
                        for (String prop : props) {
                            stmt3.setFloat(j, r.getFloat(prop));
                            j++;
                        }
                        stmt3.execute();
                    }
                }
                // delete the [t0, t1] rows.
                stmt4.setInt(1, nodeId);
                stmt4.setInt(2, tx.getStartTime());
                stmt4.setInt(3, tx.getEndTime());
                stmt4.execute();
                // insert into table from t0 with the new value.
                stmt3.setInt(1, tx.getStartTime());
                stmt3.setInt(2, nodeId);
                for(int i=3; i<props.length+3; i++){
                    stmt3.setFloat(i, tx.getValues().get(i-3));
                }
                stmt3.execute();
                conn.commit();
            } catch (SQLException | NullPointerException e) {
                System.err.println(e.getMessage());
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            return new AbstractTransaction.Result();
        };
    }

    private Req execute(SnapshotQueryTx tx) {
        return conn -> {
            List<Pair<Integer, Float>> res = new ArrayList<>();
            try (PreparedStatement stmt1 = conn.prepareStatement("select ts.node_id, " + tx.getPropertyName() + " as property from temporal_status as ts, " +
                         "(select node_id, max(t) as max_t from temporal_status where t <= ? group by node_id) as tmp where ts.node_id = tmp.node_id and ts.t = tmp.max_t")) {
                conn.setAutoCommit(false);
                int snapshotTime = tx.getTimestamp();
                stmt1.setInt(1, snapshotTime);
                ResultSet rs = stmt1.executeQuery();
                while (rs.next()) {
                    res.add(Pair.of(rs.getInt("node_id"), rs.getFloat("property")));
                }
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotQueryTx.Result result = new SnapshotQueryTx.Result();
            result.setStatus(res);
            return result;
        };
    }

    private Req execute(SnapshotAggrMaxTx tx) {
        return conn -> {
            List<Pair<Integer, Float>> res = new ArrayList<>();
            try (PreparedStatement stmt2 = conn.prepareStatement("select ts.node_id, max(" + tx.getP() + ") as max_p from temporal_status as ts, " +
                    "(select node_id, max(t) as max_t from temporal_status where t<=? group by node_id) as tmp " +
                    "where tmp.node_id=ts.node_id and tmp.max_t<=ts.t and ts.t<=? group by ts.node_id")) {
                conn.setAutoCommit(false);
                stmt2.setInt(1, tx.getT0());
                stmt2.setInt(2, tx.getT1());
                ResultSet rs = stmt2.executeQuery();
                while (rs.next()) {
                    res.add(Pair.of(rs.getInt("node_id"), rs.getFloat("max_p")));
                }
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotAggrMaxTx.Result result = new SnapshotAggrMaxTx.Result();
            result.setPropMaxValue(res);
            return result;
        };
    }

    private Req execute(SnapshotAggrDurationTx tx) {
        final TreeSet<Float> intStart = tx.getIntStartTreeSet();
        return conn -> {
            List<Triple<Integer, Float, Integer>> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select node_id, max(t) as real_st from temporal_status where t <= ? group by node_id");
                 PreparedStatement stmt1 = conn.prepareStatement("select t, " + tx.getP() + " as property from temporal_status where node_id = ? and t >= ? and t <= ? order by t")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getT0());
                ResultSet rs = stmt0.executeQuery();
                while (rs.next()) {
                    int id = rs.getInt("node_id");
                    stmt1.setInt(1, id);
                    stmt1.setInt(2, rs.getInt("real_st"));
                    stmt1.setInt(3, tx.getT1());
                    ResultSet r = stmt1.executeQuery();
                    int lastTime = -1;
                    float lastStatus = Float.MIN_VALUE;
                    HashMap<Float, Integer> buffer = new HashMap<>();
                    while (r.next()) {
                        float property = r.getFloat("property");
                        int t = r.getInt("t");
                        if (lastTime != -1) {
                            int duration = t - lastTime;
                            Float statusGrp = intStart.floor( lastStatus);
                            if(statusGrp!=null) buffer.merge(statusGrp, duration, Integer::sum);
                            lastTime = t;
                        } else {
                            lastTime = tx.getT0();
                        }
                        lastStatus = property;
                    }
                    if (lastTime != -1) {
                        Float statusGrp = intStart.floor(lastStatus);
                        if(statusGrp!=null) buffer.merge(statusGrp, tx.getT1() - lastTime + 1, Integer::sum);
                    }
                    for (Map.Entry<Float, Integer> e : buffer.entrySet()) {
                        res.add(Triple.of(id, e.getKey(), e.getValue()));
                    }
//                        System.out.println(res.size());
                }
            } catch (SQLException e) {
                e.printStackTrace();
                throw e;
            } finally {
                conn.setAutoCommit(true);
            }
            SnapshotAggrDurationTx.Result result = new SnapshotAggrDurationTx.Result();
            result.setStatusDuration(res);
            return result;
        };
    }

    private Req execute(EntityHistoryTx tx) {
        return conn -> {
            List<Triple<Integer, Integer, Float>> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select max(t) as real_st from temporal_status where t <= ? and node_id=?");
                 PreparedStatement stmt1 = conn.prepareStatement("select t, " + tx.getProp() + " as property from temporal_status where node_id = ? and t >= ? and t <= ? order by t")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getBeginTime());
                stmt0.setInt(2, tx.getEntity());
                ResultSet rs = stmt0.executeQuery();
                if (rs.next()) {
                    stmt1.setInt(1, tx.getEntity());
                    stmt1.setInt(2, rs.getInt("real_st"));
                    stmt1.setInt(3, tx.getEndTime());
                    ResultSet r = stmt1.executeQuery();
                    int lastTime = -1;
                    float lastStatus = Float.MIN_VALUE;
                    while (r.next()) {
                        float property = r.getFloat("property");
                        int t = r.getInt("t");
                        if (lastTime != -1) {
                            res.add(Triple.of(lastTime, t-1, lastStatus));
                            lastTime = t;
                        } else {
                            lastTime = tx.getBeginTime();
                        }
                        lastStatus = property;
                    }
                    if (lastTime != -1) {
                        res.add(Triple.of(lastTime, tx.getEndTime(), lastStatus));
                    }
                }
            } catch (SQLException e) {
                e.printStackTrace();
                throw e;
            } finally {
                conn.setAutoCommit(true);
            }
            EntityHistoryTx.Result result = new EntityHistoryTx.Result();
            result.setHistory(res);
            return result;
        };
    }

    private Req execute(EntityTemporalConditionTx tx) {
        return conn -> {
            List<Integer> res = new ArrayList<>();
            try (PreparedStatement stmt0 = conn.prepareStatement("select distinct ts.node_id from temporal_status as ts, " +
                    "(select node_id, max(t) as max_t from temporal_status where t<=? group by node_id) as tmp " +
                    "WHERE tmp.node_id=ts.node_id AND ts."  + tx.getP() + " >= ? AND ts." + tx.getP() + " <= ? AND "+
                    "tmp.max_t<=ts.t AND ts.t<=?")) {
                conn.setAutoCommit(false);
                stmt0.setInt(1, tx.getT0());
                stmt0.setFloat(2, tx.getVMin());
                stmt0.setFloat(3, tx.getVMax());
                stmt0.setInt(4, tx.getT1());
                ResultSet rs = stmt0.executeQuery();
                while (rs.next()) {
                    res.add(rs.getInt("node_id"));
                }
            } catch (SQLException e) {
                e.printStackTrace();
            } finally {
                conn.setAutoCommit(true);
            }
            EntityTemporalConditionTx.Result result = new EntityTemporalConditionTx.Result();
            result.setEntities(res);
            return result;
        };
    }
}
